#include "myaruco.h"

namespace myaruco {
    void readCameraParameters(cv::Mat &cameraMatrix, cv::Mat &distCoeffs)
    {
        cameraMatrix = (cv::Mat_<float>(3, 3) <<
            307.10092,   0.     , 314.4606 ,
            0.     , 307.11747, 199.90106,
            0.     ,   0.     ,   1.     );
        distCoeffs = (cv::Mat_<float>(5, 1) << -0.019218, 0.012827, -0.000142, -0.000708, 0.000000);
    }

    void maker_create()
    {
        //创建aruco标记
        cv::Mat marker;
        auto dictionary = cv::aruco::getPredefinedDictionary(cv::aruco::PREDEFINED_DICTIONARY_NAME::DICT_6X6_250);
        int i = 20;
        cv::aruco::drawMarker(dictionary, i, 200, marker, 1);
        std::string windowName = "marker" + std::to_string(i);
        imshow(windowName, marker);
        std::string path = windowName + ".jpg";
        imwrite(path, marker);
    }

    void maker_detect_image(cv::Mat test_image)
    {
        cv::Mat marker;
        cv::Ptr<cv::aruco::Dictionary> dictionary = cv::aruco::getPredefinedDictionary(cv::aruco::PREDEFINED_DICTIONARY_NAME::DICT_5X5_50);
        cv::Ptr<cv::aruco::DetectorParameters> parameters = cv::aruco::DetectorParameters::create();
        parameters->cornerRefinementMethod = cv::aruco::CORNER_REFINE_NONE;

        cv::Mat imageCopy;

        test_image.copyTo(imageCopy);
        std::vector<int> ids;
        std::vector<std::vector<cv::Point2f>>corners, rejectedImgPoints;
        cv::aruco::detectMarkers(test_image, dictionary, corners, ids, parameters, rejectedImgPoints);
        
        if (ids.size() > 0) 
        {
            cv::aruco::drawDetectedMarkers(imageCopy, corners, ids, cv::Scalar(0, 255, 0));
            cv::imshow("test", imageCopy);
            cv::waitKey(0);
        }
        else
        {
            printf("num is %ld\n",ids.size());
        }
    }

    void maker_init_dic(cv::Ptr<cv::aruco::Dictionary> &dictionary, cv::aruco::PREDEFINED_DICTIONARY_NAME dict)
    {
        // cv::Ptr<cv::aruco::Dictionary> dictionary = cv::aruco::getPredefinedDictionary(cv::aruco::DICT_5X5_50);
        dictionary = cv::aruco::getPredefinedDictionary(dict);
    }

    void maker_detect_video(cv::Mat image, cv::Ptr<cv::aruco::Dictionary> &dictionary)
    {
        cv::Mat imageCopy;
        image.copyTo(imageCopy);
        std::vector<int> ids;
        std::vector<std::vector<cv::Point2f>> corners;
        cv::aruco::detectMarkers(image, dictionary, corners, ids);
        // if at least one marker detected
        if (ids.size() > 0)
            cv::aruco::drawDetectedMarkers(imageCopy, corners, ids);
        // else
        //     printf("num is %ld\n",ids.size());
        cv::imshow("out", imageCopy);
    }

    void maker_pose_estimation(cv::Mat image, cv::Ptr<cv::aruco::Dictionary> &dictionary, cv::Mat &cameraMatrix, cv::Mat &distCoeffs)
    {
        // camera parameters are read from somewhere
        cv::Mat imageCopy;
        image.copyTo(imageCopy);
        std::vector<int> ids;
        std::vector<std::vector<cv::Point2f>> corners;
        cv::aruco::detectMarkers(image, dictionary, corners, ids);
        // if at least one marker detected
        if (ids.size() > 0) {
            cv::aruco::drawDetectedMarkers(imageCopy, corners, ids);
            std::vector<cv::Vec3d> rvecs, tvecs;
            cv::aruco::estimatePoseSingleMarkers(corners, 0.05, cameraMatrix, distCoeffs, rvecs, tvecs);
            // draw axis for each marker
            for(int i=0; i<ids.size(); i++)
                cv::aruco::drawAxis(imageCopy, cameraMatrix, distCoeffs, rvecs[i], tvecs[i], 0.1);
        }
        cv::imshow("out", imageCopy);
    }
}